# ai_detection/yolo_detector.py - YOLOv8 object detection

import os
import cv2
import numpy as np
import logging
from ultralytics import YOLO

logger = logging.getLogger(__name__)


class YoloDetector:
    """
    Class for detecting objects in frames using YOLOv8
    """

    def __init__(self, model_name="yolov8n.pt"):
        self.model_path = model_name
        self.model = None
        self.initialized = False

        # Define classes we're interested in
        self.football_classes = {
            0: 'person',  # Players, referees, staff
            32: 'sports ball',  # Football
            37: 'sports',  # General sports equipment
            38: 'fan',  # Might detect fans
            41: 'cup',  # Trophy
        }

        # Try to initialize the model
        self.initialize()

    def initialize(self):
        """Initialize the YOLO model"""
        try:
            logger.info(f"Initializing YOLOv8 detector with model: {self.model_path}")
            self.model = YOLO(self.model_path)
            self.initialized = True
            logger.info("YOLOv8 detector initialized successfully")
        except Exception as e:
            logger.error(f"Failed to initialize YOLOv8 detector: {str(e)}")
            self.initialized = False

    def detect_objects(self, frame):
        """
        Detect objects in a single frame

        Args:
            frame: BGR image (numpy array)

        Returns:
            List of detection objects [class_id, confidence, [x1, y1, x2, y2]]
        """
        if not self.initialized:
            logger.warning("YOLOv8 detector not initialized")
            return []

        try:
            # Run detection
            results = self.model(frame, verbose=False)

            # Process results - convert to standard format
            detections = []
            for r in results:
                boxes = r.boxes
                for box in boxes:
                    # Get class_id
                    cls_id = int(box.cls.item())

                    # Skip if not in our football classes
                    if cls_id not in self.football_classes:
                        continue

                    # Get confidence
                    conf = box.conf.item()

                    # Get coordinates (convert to [x1,y1,x2,y2] format)
                    xyxy = box.xyxy.cpu().numpy()[0]

                    # Add to detections
                    detections.append({
                        'class_id': cls_id,
                        'class_name': self.football_classes.get(cls_id, 'unknown'),
                        'confidence': conf,
                        'bbox': xyxy.tolist()
                    })

            return detections

        except Exception as e:
            logger.error(f"Error in object detection: {str(e)}")
            return []

    def detect_players(self, frame, min_confidence=0.4):
        """
        Detect football players in the frame

        Args:
            frame: BGR image
            min_confidence: Minimum confidence threshold

        Returns:
            List of player bounding boxes
        """
        detections = self.detect_objects(frame)

        # Filter for people (class 0) with high confidence
        players = [d for d in detections
                   if d['class_id'] == 0 and d['confidence'] >= min_confidence]

        return players

    def detect_ball(self, frame, min_confidence=0.4):
        """
        Detect football in the frame

        Args:
            frame: BGR image
            min_confidence: Minimum confidence threshold

        Returns:
            Ball detection if found, None otherwise
        """
        detections = self.detect_objects(frame)

        # Filter for sports ball (class 32) with high confidence
        balls = [d for d in detections
                 if d['class_id'] == 32 and d['confidence'] >= min_confidence]

        # Return the highest confidence ball
        if balls:
            return max(balls, key=lambda x: x['confidence'])

        return None

    def detect_goal(self, frames):
        """
        Detect if a sequence of frames contains a goal

        Args:
            frames: List of frames

        Returns:
            Confidence score that this is a goal (0-1)
        """
        goal_confidence = 0.0

        # A goal typically involves:
        # 1. Ball moving towards goal area
        # 2. Players celebrating
        # 3. Sudden camera movement (cut to celebration)

        try:
            ball_positions = []
            player_counts = []

            # Analyze each frame
            for frame in frames:
                # Detect ball
                ball = self.detect_ball(frame, min_confidence=0.3)
                if ball:
                    ball_positions.append(ball['bbox'])
                else:
                    ball_positions.append(None)

                # Count players
                players = self.detect_players(frame, min_confidence=0.3)
                player_counts.append(len(players))

            # Analyze ball trajectory
            if len(ball_positions) >= 3:
                # Check if ball is moving towards goal (simplified)
                valid_positions = [pos for pos in ball_positions if pos is not None]

                if len(valid_positions) >= 2:
                    # Check horizontal movement patterns
                    x_positions = [pos[0] for pos in valid_positions]
                    x_movement = max(x_positions) - min(x_positions)

                    # Check vertical movement (ball going down)
                    y_positions = [pos[1] for pos in valid_positions]

                    if x_movement > 50:  # Significant horizontal movement
                        goal_confidence += 0.3

                    # Check for player count changes (celebration)
                    if player_counts:
                        avg_players = sum(player_counts) / len(player_counts)
                        max_players = max(player_counts)

                        # More players appearing could indicate celebration
                        if max_players > avg_players * 1.5:
                            goal_confidence += 0.3

            # Limit confidence to 0-1 range
            goal_confidence = min(1.0, goal_confidence)

        except Exception as e:
            logger.error(f"Error in goal detection: {str(e)}")

        return goal_confidence

    def detect_celebration(self, frames):
        """
        Detect if a sequence of frames contains celebration

        Args:
            frames: List of frames

        Returns:
            Confidence score that this is a celebration (0-1)
        """
        celebration_confidence = 0.0

        try:
            player_counts = []
            player_positions = []

            # Analyze each frame
            for frame in frames:
                # Detect players
                players = self.detect_players(frame, min_confidence=0.3)
                player_counts.append(len(players))

                # Store player positions
                frame_positions = [player['bbox'] for player in players]
                player_positions.append(frame_positions)

            if player_counts:
                # Celebrations typically have:
                # 1. Players grouped together
                # 2. Players with raised arms
                # 3. Sometimes fewer players visible as camera focuses on celebrating players

                # Check for player clustering
                for positions in player_positions:
                    if len(positions) >= 2:
                        # Calculate average distance between players
                        distances = []
                        for i in range(len(positions)):
                            for j in range(i + 1, len(positions)):
                                # Center point of each bbox
                                center1 = [(positions[i][0] + positions[i][2]) / 2,
                                           (positions[i][1] + positions[i][3]) / 2]
                                center2 = [(positions[j][0] + positions[j][2]) / 2,
                                           (positions[j][1] + positions[j][3]) / 2]

                                # Euclidean distance
                                dist = ((center1[0] - center2[0]) ** 2 +
                                        (center1[1] - center2[1]) ** 2) ** 0.5
                                distances.append(dist)

                        # Low average distance indicates clustering (celebration)
                        if distances and sum(distances) / len(distances) < 200:
                            celebration_confidence += 0.4
                            break

                # Check for player count consistency (focused shots)
                std_dev = np.std(player_counts)
                if std_dev < 2 and np.mean(player_counts) > 2:
                    celebration_confidence += 0.3

                # Check for player count increase (teammates joining celebration)
                if len(player_counts) > 2:
                    if max(player_counts) > min(player_counts) * 1.5:
                        celebration_confidence += 0.3

            # Limit confidence to 0-1 range
            celebration_confidence = min(1.0, celebration_confidence)

        except Exception as e:
            logger.error(f"Error in celebration detection: {str(e)}")

        return celebration_confidence